"""
Author: Morphlng
Date: 2024-03-19 13:52:27
LastEditTime: 2024-03-19 13:52:28
LastEditors: Morphlng
Description: Test for rotation
FilePath: /DrivingGym/tests/common/geom/rotation_test.py
"""

import math

import numpy as np

from driving_gym.simulation.common.geom import Rotation
from driving_gym.simulation.common.stub import try_import_carla

carla = try_import_carla()


class TestRotation:

    def test_initialization(self):
        rotation = Rotation(30, 45, 60)
        assert (
            rotation.pitch == 30 and rotation.yaw == 45 and rotation.roll == 60
        ), "Initialization of Rotation failed."

    def test_get_right_vector(self):
        random_rotation = np.random.rand(3) * 360
        rotation = Rotation(*random_rotation)
        right_vec = rotation.get_right_vector()

        carla_rotation = carla.Rotation(*random_rotation)
        carla_right_vec = carla_rotation.get_right_vector()

        assert np.allclose(
            [right_vec.x, right_vec.y, right_vec.z],
            [carla_right_vec.x, carla_right_vec.y, carla_right_vec.z],
            atol=1e-5,
        ), "Right vector calculation failed."

    def test_get_forward_vector(self):
        random_rotation = np.random.rand(3) * 360
        rotation = Rotation(*random_rotation)
        forward_vec = rotation.get_forward_vector()

        carla_rotation = carla.Rotation(*random_rotation)
        carla_forward_vec = carla_rotation.get_forward_vector()

        assert np.allclose(
            [forward_vec.x, forward_vec.y, forward_vec.z],
            [carla_forward_vec.x, carla_forward_vec.y, carla_forward_vec.z],
            atol=1e-5,
        ), "Forward vector calculation failed."

    def test_get_up_vector(self):
        random_rotation = np.random.rand(3) * 360
        rotation = Rotation(*random_rotation)
        up_vec = rotation.get_up_vector()

        carla_rotation = carla.Rotation(*random_rotation)
        carla_up_vec = carla_rotation.get_up_vector()

        assert np.allclose(
            [up_vec.x, up_vec.y, up_vec.z],
            [carla_up_vec.x, carla_up_vec.y, carla_up_vec.z],
            atol=1e-5,
        ), "Up vector calculation failed."

    def test_as_numpy_array(self):
        rotation = Rotation(30, 45, 60)
        assert np.array_equal(
            rotation.as_numpy_array(), np.array([30, 45, 60])
        ), "Numpy array conversion failed."

    def test_copy(self):
        rotation = Rotation(30, 45, 60)
        rotation_copy = rotation.copy()
        assert (
            rotation_copy == rotation and rotation_copy is not rotation
        ), "Copy of Rotation failed."

    def test_eq(self):
        rotation1 = Rotation(30, 45, 60)
        rotation2 = Rotation(30, 45, 60)
        assert rotation1 == rotation2, "Equality check failed."

    def test_ne(self):
        rotation1 = Rotation(30, 45, 60)
        rotation2 = Rotation(60, 45, 30)
        assert rotation1 != rotation2, "Inequality check failed."
